#include "sample.h"

#if (defined __GNUC__) && (__GNUC__>4 || __GNUC_MINOR__>=7)
  #undef _GLIBCXX_ATOMIC_BUILTINS
  #undef _GLIBCXX_USE_INT128
#endif

#include <iostream>
#define EIGEN_USE_THREADS
#define EIGEN_USE_GPU
// #define EIGEN_GPUCC
#define EIGEN_TEST_NO_LONGDOUBLE
#define EIGEN_TEST_NO_COMPLEX
#define EIGEN_TEST_FUNC cxx11_tensor_cuda
#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
#include <stdio.h>
#include <unsupported/Eigen/CXX11/Tensor>

#define N 8
using namespace Eigen;
typedef Eigen::GpuStreamDevice GPUStream;
typedef Eigen::GpuDevice GPUDevice;
typedef Eigen::DenseIndex IndexType;
typedef Eigen::ThreadPoolDevice CPUDevice;
#define NDIMS 3
template<typename T>
using TTensor =  Eigen::TensorMap<Eigen::Tensor<T, NDIMS, Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using Flat = Eigen::TensorMap<Eigen::Tensor<T, 1, Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using  ConstScalar = Eigen::TensorMap<Eigen::TensorFixedSize<const T, Eigen::Sizes<>,Eigen::RowMajor, IndexType>,Eigen::Aligned>;
template<typename T>
using ConstFlat =  Eigen::TensorMap<
Eigen::Tensor<const T, 1, Eigen::RowMajor, IndexType>, Eigen::Aligned>;
const int esp = 1e5 + 1;

// Eigen::internal::TensorExecutor<Eigen::TensorAssignOp<Eigen::TensorMap<Eigen::Tensor<float, 1, 1, int>, 16, Eigen::MakePointer>, Eigen::TensorMap<Eigen::Tensor<float, 1, 1, int>, 16, Eigen::MakePointer> const> const, Eigen::GpuDevice, true, (Eigen::internal::TiledEvaluation)0>::run(
//   Eigen::TensorAssignOp<Eigen::TensorMap<Eigen::Tensor<float, 1, 1, int>, 16, Eigen::MakePointer>, Eigen::TensorMap<Eigen::Tensor<float, 1, 1, int>, 16, Eigen::MakePointer> const> const&, Eigen::GpuDevice const&)
template<typename T>
struct test{
  void call(T* _m,T* _v,T* _var,T* _grad,T* f1, T*f2,T*f3,T*f4,T*f5) {
    GPUStream stream;
    GPUDevice d(&stream, 256);
    Flat<T> m(_m,N),v(_v,N),var(_var,N);
    ConstScalar<T> beta1_power(f1), lr(f2),beta1(f3),beta2(f4),epsilon(f5);
    ConstFlat<T> grad(_grad,N);
    Eigen::array<typename TTensor<T>::Index, 1> bcast;
    bcast[0] = grad.dimension(0);
    Eigen::Sizes<1> single;
    const auto one = static_cast<T>(1.0);
    m.device(d) =
      m + (beta1.constant(one) - beta1).reshape(single).broadcast(bcast) *
      (grad - m);
    v.device(d) =
      (beta2.reshape(single).broadcast(bcast) * v).cwiseMax(grad.abs());
    var.device(d) -=
      (lr / (beta1_power.constant(one) -
          beta1_power)).reshape(single).broadcast(bcast) *
      (m / (v + epsilon.reshape(single).broadcast(bcast)));
  }
};
int eigen_gpu_sample() {
  float _v[] = {0.3042 ,0.525  ,0.432  ,0.2913 ,0.612  ,0.1395 ,0.2922 ,0.3665};
  float _m[] = {0.601   ,0.708   ,0.02058 ,0.9697  ,0.8325  ,0.2123  ,0.1819  ,0.1833};
  float _grad[] = {0.456   ,0.785   ,0.1997  ,0.514   ,0.5923  ,0.04645 ,0.6074  ,0.1705 };
  float _var[] = {0.3745  ,0.9507  ,0.732   ,0.5986  ,0.156   ,0.156   ,0.05807 ,0.866 };
  float beta1_power = 0.9;
  float lr = 0.01;
  float beta1 = 0.9;
  float beta2 = 0.99;
  float epsilon = 1e-8;
  float *_m_dev, *_v_dev, *_var_dev, *_grad_dev, *beta1_power_dev,
        *lr_dev, *beta1_dev, *beta2_dev, *epsilon_dev;
  cudaMalloc((float**)(&_m_dev), sizeof(float)*N);
  cudaMalloc((float**)(&_v_dev), sizeof(float)*N);
  cudaMalloc((float**)(&_var_dev), sizeof(float)*N);
  cudaMalloc((float**)(&_grad_dev), sizeof(float)*N);
  cudaMalloc((float**)(&beta1_power_dev), sizeof(float));
  cudaMalloc((float**)(&beta1_dev), sizeof(float));
  cudaMalloc((float**)(&beta2_dev), sizeof(float));
  cudaMalloc((float**)(&epsilon_dev), sizeof(float));
  cudaMalloc((float**)(&lr_dev), sizeof(float));
  cudaMemcpy(_m_dev, _m, sizeof(float)*N, cudaMemcpyHostToDevice);
  cudaMemcpy(_v_dev, _v, sizeof(float)*N, cudaMemcpyHostToDevice);
  cudaMemcpy(_var_dev, _var, sizeof(float)*N, cudaMemcpyHostToDevice);
  cudaMemcpy(_grad_dev, _grad, sizeof(float)*N, cudaMemcpyHostToDevice);
  cudaMemcpy(lr_dev, &lr, sizeof(float), cudaMemcpyHostToDevice);
  cudaMemcpy(beta1_power_dev, &beta1_power, sizeof(float), cudaMemcpyHostToDevice);
  cudaMemcpy(beta1_dev, &beta1, sizeof(float), cudaMemcpyHostToDevice);
  cudaMemcpy(beta2_dev, &beta2, sizeof(float), cudaMemcpyHostToDevice);
  cudaMemcpy(epsilon_dev, &epsilon, sizeof(float), cudaMemcpyHostToDevice);

  test<float> gpu_gen;
  gpu_gen.call( _m_dev, _v_dev, _var_dev, _grad_dev,
      beta1_power_dev, lr_dev, beta1_dev, beta2_dev, epsilon_dev);
  float ans[N];
  cudaMemcpy(ans, _var_dev, sizeof(float)*N, cudaMemcpyDeviceToHost);
  cudaDeviceSynchronize();
  printf("gpu: %f\n", ans[0]);
  cudaFree(_m_dev);
  cudaFree(_v_dev);
  cudaFree(_var_dev);
  cudaFree(_grad_dev);
  cudaFree(beta1_power_dev);
  cudaFree(lr_dev);
  cudaFree(beta1_dev);
  cudaFree(beta2_dev);
  cudaFree(epsilon_dev);
  // printf("CUDA_ERROR : %s\n",cudaGetErrorString(cudaGetLastError()));
  return 0;
}
